from mpl_toolkits import mplot3d
import numpy as np
import matplotlib.pyplot as plt
from python_ai.ML.lin_regression.xlib import *
from scipy.optimize import curve_fit

np.random.seed(1)
dataOri = np.loadtxt(r'../../../data/ex1data2.txt', delimiter=',')
m = dataOri.shape[0]
# data = dataOri
data, mu, sigma = scale_feature_data(dataOri)
print(m)
print(data[:5])
x = data[:, 0].reshape(m, 1)
y = data[:, 1].reshape(m, 1)
z = data[:, 2].reshape(m, 1)
print(x[:5])
print(y[:5])
print(z[:5])

XarrAx = np.linspace(data[:, 0].min(), data[:, 0].max(), m)
YarrAx = np.linspace(data[:, 1].min(), data[:, 1].max(), m)

X = np.c_[np.ones([m, 1]), x, y]
print(X[:5])
print(x.shape, y.shape, z.shape, X.shape)

num_iters = 15000
theata, history, xscores = gradient_descent_algorithm(X, z, alpher=0.001, num_iters=num_iters)
print('THETA', theata)


def func_fit(X, theata):
    # Y = theata[0] * X[:, 0] + theata[1] * X[:, 1] + theata[2] * X[:, 2]
    Y = np.dot(X, theata)
    return Y


def func_fit4curve_fit(X, a, b, c):
    Y = a * X[:, 0] + b * X[:, 1] + c * X[:, 2]
    return Y


def on_ax(ax):
    ax.set_title('梯度下降算法')
    ax.scatter3D(x, y, z)
    Xy = np.c_[np.ones([m, 1]), XarrAx.reshape(m, 1), YarrAx.reshape(m, 1)]
    ax.plot3D(XarrAx, YarrAx, func_fit(Xy, theata).ravel(), 'r-', label='Gradient Descent Regression')
    popt, popv = curve_fit(func_fit4curve_fit, X, z.ravel())
    ax.plot3D(XarrAx, YarrAx, func_fit4curve_fit(Xy, popt[0], popt[1], popt[2]), '*', label='scipy.optimize.curve_fit')
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_zlabel('z')
    ax.legend()


plt.ioff()

fig = plt.figure(figsize=[12, 8])
pr = 1
pc = 1
plt.rcParams['font.sans-serif'] = ['Simhei']  # 设置中文字体为黑体
plt.rcParams['axes.unicode_minus'] = False  # 显示负号

ax = fig.add_subplot(pr, pc, 1, projection='3d')
on_ax(ax)

fig = plt.figure(figsize=[12, 8])
pr = 2
pc = 2
plt.rcParams['font.sans-serif'] = ['Simhei']  # 设置中文字体为黑体
plt.rcParams['axes.unicode_minus'] = False  # 显示负号

ax = fig.add_subplot(pr, pc, 1, projection='3d')
on_ax(ax)

plt.subplot(pr, pc, 2)
plt.title('Cost function value')
xx = range(num_iters)
plt.plot(xx, history)
plt.grid()

plt.subplot(pr, pc, 3)
plt.title('Cost function value, 2nd half')
xx = range(num_iters // 2, num_iters)
plt.plot(xx, history[num_iters // 2:])
plt.grid()

if len(xscores) > 0:
    print(xscores[:10])
    print(xscores[-10:])
    plt.subplot(pr, pc, 4)
    plt.title('Score function value')
    xx = range(num_iters)
    plt.plot(xx, xscores)
    plt.grid()

plt.show()
